import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from sklearn.datasets import load_diabetes  # 糖尿病数据集
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import ElasticNet, ElasticNetCV
from sklearn.metrics import mean_squared_error, r2_score

# 设置 Matplotlib 中文显示（可选）
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # Windows 使用 SimHei
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 1. 加载数据集
data = load_diabetes()
X, y = data.data, data.target  # X 是特征矩阵，y 是目标变量

# 2. 拆分训练集和测试集（80% 训练，20% 测试）
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 3. 数据标准化（弹性网络对特征缩放敏感）
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# 4. 交叉验证寻找最佳 α（正则化强度）和 l1_ratio（L1/L2 比例）
elastic_cv = ElasticNetCV(l1_ratio=np.linspace(0.1, 1, 10),
                          alphas=np.logspace(-3, 1, 10),
                          cv=5, random_state=42)
elastic_cv.fit(X_train, y_train)

best_alpha = elastic_cv.alpha_
best_l1_ratio = elastic_cv.l1_ratio_
print(f"最佳 α 值: {best_alpha:.4f}")
print(f"最佳 L1 比例: {best_l1_ratio:.4f}")

# 5. 训练最终的弹性网络模型
elastic_net = ElasticNet(alpha=best_alpha, l1_ratio=best_l1_ratio)
elastic_net.fit(X_train, y_train)

# 6. 预测
y_pred = elastic_net.predict(X_test)

# 7. 计算评价指标
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)
print(f"弹性网络回归的 MSE: {mse:.4f}")
print(f"弹性网络回归的 R² 分数: {r2:.4f}")

# 8. 数据可视化

# (1) 绘制特征重要性（弹性网络回归系数）
plt.figure(figsize=(10, 5))
plt.bar(range(X.shape[1]), elastic_net.coef_)
plt.xlabel("特征索引")
plt.ylabel("回归系数")
plt.title("弹性网络回归：特征重要性")
plt.xticks(range(X.shape[1]), data.feature_names, rotation=45)
plt.show()

# (2) 预测值 vs 真实值
plt.figure(figsize=(8, 6))
plt.scatter(y_test, y_pred, color="blue", alpha=0.6, label="预测值")
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], "r--", label="理想拟合")
plt.xlabel("真实值")
plt.ylabel("预测值")
plt.title("弹性网络回归：预测值 vs 真实值")
plt.legend()
plt.grid()
plt.show()

# (3) 误差分布直方图
errors = y_test - y_pred
plt.figure(figsize=(8, 5))
plt.hist(errors, bins=20, color="purple", alpha=0.7, edgecolor="black")
plt.xlabel("误差（真实值 - 预测值）")
plt.ylabel("频数")
plt.title("弹性网络回归误差分布")
plt.grid()
plt.show()
